Rebase PR #1479 onto current main (post-PRs #1507/#1526/#1527/#1529).
The original E2E fix's `pad_batch` change is dropped since #1529's
prompt-based mini-batch boundaries removed the need to pad to
`mini_batch_size * n_samples`.

- merge_stepwise_output() in trainer_utils.py collapses multi-turn
  step-wise GeneratorOutput sequences into single sequences when
  consecutive turns share a common prefix, reducing training cost
  from O(T^2) to O(T).
- trainer.py: call merge before extracting generator fields, update
  uids from merged trajectory_ids, emit generate/num_seq_{before,after}_merge.
- Add generator.merge_stepwise_output config flag.
- run_search.sh: MERGE_STEPWISE env var.
- 16 CPU-only tests covering all 3 merging cases, partial merges, and
  validation asserts.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…1532)

## Summary

This PR implements prefix-aware merging for step-wise training, guarded
by a flag `cfg.generator.merge_stepwise_output` that defaults to False.
During step-wise training, within a trajectory, when consecutive steps
share the same prefix (i.e. no re-tokenization drift or context
management like thinking token stripping), we collapse into a single
`GeneratorOutput` entry. This can reduce the O(T²) training cost
introduced by step-wise (T being number of turns).

- `merge_stepwise_output()` in `skyrl/train/utils/trainer_utils.py`
implements greedy merging: for consecutive turns in the same trajectory
where `prompt[i] + response[i]` is a prefix of `prompt[i+1]`, merge into
one entry. Response tokens concatenated with the observation-delta
(loss-masked to 0) between turns; per-token fields (`loss_masks`,
`rewards`, `rollout_logprobs`) align accordingly; per-turn fields
(`stop_reason`, `is_last_step`, `trajectory_id`) take the last turn's
value.
- `RayPPOTrainer.postprocess_generator_output` calls
`merge_stepwise_output` when `generator.merge_stepwise_output=true`,
updates `uids` from the merged `trajectory_ids`, and logs
`generate/num_seq_{before,after}_merge`.
- Since `uids` may need to be modified, update the signature of
`postprocess_generator_output` to return both `generator_output` and
`uids`, changing various caller places
- New `generator.merge_stepwise_output` config flag (default false).
- `examples/train/search/run_search.sh` accepts `MERGE_STEPWISE=true`
env var to pass the flag through.
- 16 CPU-only unit tests in `tests/train/test_merge_stepwise_output.py`
cover the three merge cases, partial merges, prefix mismatches,
single-turn passthrough, per-trajectory scalar rewards, and
required-field asserts.

## Test plan

- [x] `pytest tests/train/test_merge_stepwise_output.py` — 16 passed
- [x] `pytest tests/train/test_trainer_utils.py
tests/train/test_prompt_mini_batch.py` — 58 passed (existing tests
unaffected)
- [x] E2E: Search-R1 step-wise GRPO run on Qwen2.5-3B-Instruct, 8×H100,
`MERGE_STEPWISE=true`.

### Curves

With pricesly the same setup as #1529 , we do:

```bash
MERGE_STEPWISE=true USE_CONVERSATION_MULTI_TURN=true STEP_WISE=true bash examples/train/search/run_search.sh \
  generator.inference_engine.num_engines=8 \
  generator.inference_engine.tensor_parallel_size=1
```
See PR description for more.

Co-authored-by: Deep Sheth
[deepsheth3@users.noreply.github.com](mailto:deepsheth3@users.noreply.github.com)

…1538)

## Summary

This PR implements prefix-aware merging for step-wise training, guarded
by a flag `cfg.generator.merge_stepwise_output` that defaults to False.
During step-wise training, within a trajectory, when consecutive steps
share the same prefix (i.e. no re-tokenization drift or context
management like thinking token stripping), we collapse into a single
`GeneratorOutput` entry. This can reduce the O(T²) training cost
introduced by step-wise (T being number of turns).

- `merge_stepwise_output()` in `skyrl/train/utils/trainer_utils.py`
implements greedy merging: for consecutive turns in the same trajectory
where `prompt[i] + response[i]` is a prefix of `prompt[i+1]`, merge into
one entry. Response tokens concatenated with the observation-delta
(loss-masked to 0) between turns; per-token fields (`loss_masks`,
`rewards`, `rollout_logprobs`) align accordingly; per-turn fields
(`stop_reason`, `is_last_step`, `trajectory_id`) take the last turn's
value.
- `RayPPOTrainer.postprocess_generator_output` calls
`merge_stepwise_output` when `generator.merge_stepwise_output=true`,
updates `uids` from the merged `trajectory_ids`, and logs
`generate/num_seq_{before,after}_merge`.
- Since `uids` may need to be modified, update the signature of
`postprocess_generator_output` to return both `generator_output` and
`uids`, changing various caller places
- New `generator.merge_stepwise_output` config flag (default false).
- `examples/train/search/run_search.sh` accepts `MERGE_STEPWISE=true`
env var to pass the flag through.
- 16 CPU-only unit tests in `tests/train/test_merge_stepwise_output.py`
cover the three merge cases, partial merges, prefix mismatches,
single-turn passthrough, per-trajectory scalar rewards, and
required-field asserts.

## Test plan

- [x] `pytest tests/train/test_merge_stepwise_output.py` — 16 passed
- [x] `pytest tests/train/test_trainer_utils.py
tests/train/test_prompt_mini_batch.py` — 58 passed (existing tests
unaffected)
- [x] E2E: Search-R1 step-wise GRPO run on Qwen2.5-3B-Instruct, 8×H100,
`MERGE_STEPWISE=true`.

### Curves

With pricesly the same setup as #1529 , we do:

```bash
MERGE_STEPWISE=true USE_CONVERSATION_MULTI_TURN=true STEP_WISE=true bash examples/train/search/run_search.sh \
  generator.inference_engine.num_engines=8 \
  generator.inference_engine.tensor_parallel_size=1
```

See PR description for more.

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Deep Sheth <deepsheth3@users.noreply.github.com>